#!/usr/bin/env python3
# utils/helpers.py
import os
import torch
from torch.utils.tensorboard import SummaryWriter

def init_tensorboard(log_dir):
    """初始化TensorBoard日志记录器"""
    os.makedirs(log_dir, exist_ok=True)
    return SummaryWriter(log_dir)

"""
@TODO
"""

def save_checkpoint(epoch, generator, discriminator, g_optimizer, d_optimizer, save_dir, logger=None):
    """保存模型检查点"""
    os.makedirs(save_dir, exist_ok=True)
    checkpoint_path = os.path.join(save_dir, f"checkpoint_epoch_{epoch}.pth")
    
    torch.save({
        'epoch': epoch,
        'generator_state_dict': generator.state_dict(),
        'discriminator_state_dict': discriminator.state_dict(),
        'g_optimizer_state_dict': g_optimizer.state_dict(),
        'd_optimizer_state_dict': d_optimizer.state_dict(),
    }, checkpoint_path)
    
    if logger:
        logger.info(f"已保存检查点至: {checkpoint_path}")
    else:
        print(f"已保存检查点至: {checkpoint_path}")